#!/usr/bin/python
#coding:utf-8
'''
将决策树结构保存为图片或pdf，pdf是矢量图。
场景:matplotlib、Graphviz
安装： Graphviz下载解压到安装路径：ttp://www.graphviz.org/download/
       matplotlib 安装：python -m pip install pydotplus

'''

import pydotplus 
#导入Graphviz，也可配入环境变量
import os 
os.environ["PATH"] += os.pathsep + 'D:/python-space/tool/release/bin'

from sklearn.datasets import load_iris
from sklearn import tree
from IPython.display import Image  
import pydotplus 


#载入sciki-learn决策
iris = load_iris()
clf = tree.DecisionTreeClassifier()
clf = clf.fit(iris.data, iris.target)

#存储模型到文件 iris.dot
# with open("iris.dot", 'w') as f:
#     f = tree.export_graphviz(clf, out_file=f)

#算法模型图生成pdf文件
# dot_data = tree.export_graphviz(clf, out_file=None) #不指定类别颜色
dot_data = tree.export_graphviz(clf, out_file=None,  #指定不同类别颜色
                         feature_names=iris.feature_names,  
                         class_names=iris.target_names,  
                         filled=True, rounded=True,  
                         special_characters=True) 

graph = pydotplus.graph_from_dot_data(dot_data) 

graph.write_png("iris.png")
# graph.write_pdf("iris.pdf")
# Image(graph.create_png()) 